from .u_net import get_unet
from .vgg import get_vgg_segmentation
from .resnet import get_resnet_segmentation

__all__ = ['get_unet', 'get_vgg_segmentation', 'get_resnet_segmentation']


def get_model(model_name, n_channels=3, n_classes=2, pretrained=False, **kwargs):
    """
    根据模型名称获取对应的模型实例
    
    Args:
        model_name (str): 模型名称，支持 'unet', 'vgg', 'resnet'
        n_channels (int): 输入通道数
        n_classes (int): 分类数
        pretrained (bool): 是否使用预训练权重
        **kwargs: 其他模型特定参数
    
    Returns:
        torch.nn.Module: 模型实例
    """
    model_name = model_name.lower()
    
    if model_name == 'unet':
        return get_unet(n_channels=n_channels, n_classes=n_classes, **kwargs)
    elif model_name == 'vgg':
        return get_vgg_segmentation(n_classes=n_classes, pretrained=pretrained, **kwargs)
    elif model_name == 'resnet':
        return get_resnet_segmentation(n_classes=n_classes, pretrained=pretrained, **kwargs)
    else:
        raise ValueError(f"Unsupported model: {model_name}. Supported models: unet, vgg, resnet")

